package smartcase.base.access;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import lombok.Data;
import org.jetbrains.annotations.NotNull;
import org.springframework.core.MethodParameter;
import org.springframework.http.MediaType;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;
import smartcase.base.ApiResponse;
import smartcase.base.log.Log;
import smartcase.base.log.LogFactory;
import smartcase.config.access.AccessConfiguration;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.LinkedHashMap;
import java.util.UUID;
import java.util.concurrent.TimeUnit;

import static java.util.Optional.ofNullable;
import static smartcase.base.ApplicationContextProviderKt.getCommonInfoFromContext;
import static smartcase.common.constants.ProjectConstants.LOGCTX_INFO_KEY;

@Component
@ControllerAdvice
@Data
public class AccessInterceptor implements HandlerInterceptor, ResponseBodyAdvice<ApiResponse> {

  private Log log = LogFactory.getLog(AccessInterceptor.class);

  private TimeUnit timeUnit;
  private Long timeWindow;
  private Integer maxAccessNumThreshold;
  private Long selfProtectionMaxAccessNum;
  private final Cache<String, Integer> overflowLimitCache;

  public AccessInterceptor(AccessConfiguration cfg) {

    this.timeUnit = ofNullable(cfg.getTimeUnit()).orElse(TimeUnit.SECONDS);
    this.maxAccessNumThreshold = ofNullable(cfg.getMaxAccessNumThreshold()).orElse(50);
    this.timeWindow = ofNullable(cfg.getTimeWindow()).orElse(3L);
    this.selfProtectionMaxAccessNum = ofNullable(cfg.getSelfProtectionMaxAccessNum()).orElse(6000L);

    this.overflowLimitCache =
        CacheBuilder.newBuilder()
            .expireAfterWrite(timeWindow, this.timeUnit)
            .softValues()
            .recordStats()
            .build();
  }

  @Override
  public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler)
      throws Exception {
    LogCtx logCtx = new LogCtx();
    logCtx.setRequestId(UUID.randomUUID().toString());
    logCtx.setPath(request.getServletPath());
    if (handler instanceof HandlerMethod) {
      HandlerMethod h = (HandlerMethod) handler;
      logCtx.setEndpoint(h.getMethod().getName());
    }
    request.setAttribute(LOGCTX_INFO_KEY, logCtx);
    String ip = request.getRemoteHost();

    // Limit Flow
    if (overflowLimitCache.stats().evictionCount() >= selfProtectionMaxAccessNum) {
      return false;
    }
    Integer accessCount = overflowLimitCache.get(ip, () -> 0);
    Integer maxAccessLimit = 20;
    if (accessCount >= maxAccessLimit) {
      return false;
    }
    overflowLimitCache.put(ip, accessCount + 1);

    return true;
  }

  @Override
  public boolean supports(@NotNull MethodParameter methodParameter, @NotNull Class type) {
    return ApiResponse.class.equals(type);
  }

  @Override
  public ApiResponse beforeBodyWrite(
      ApiResponse respBody,
      MethodParameter parameter,
      MediaType mediaType,
      Class type,
      ServerHttpRequest request,
      ServerHttpResponse response) {

    String reqId;
    LogCtx logCtx;
    if ((logCtx = getCommonInfoFromContext()) != null && (reqId = logCtx.getRequestId()) != null) {
      respBody.setRequestId(reqId);
      logCtx.setEndTime(System.currentTimeMillis());
      log.info(logCtx, "access end");
    }
    return respBody;
  }
}
